{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 解读 {class}`~tvm.ir.base.EnvFunc`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import tvm.testing # 保证 \"attrs.TestAttrs\" 被注册"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mInit signature:\u001b[39m tvm.ir.EnvFunc(self, /, *args, **kwargs)\n",
      "\u001b[31mSource:\u001b[39m        \n",
      "@register_object(\u001b[33m\"ir.EnvFunc\"\u001b[39m)\n",
      "\u001b[38;5;28;01mclass\u001b[39;00m EnvFunc(Object):\n",
      "    \u001b[33m\"\"\"Environment function.\u001b[39m\n",
      "\n",
      "\u001b[33m    This is a global function object that can be serialized by its name.\u001b[39m\n",
      "\u001b[33m    \"\"\"\u001b[39m\n",
      "\n",
      "    \u001b[38;5;28;01mdef\u001b[39;00m __call__(self, *args):\n",
      "        \u001b[38;5;28;01mreturn\u001b[39;00m _ffi_api.EnvFuncCall(self, *args)  \u001b[38;5;66;03m# type: ignore # pylint: disable=no-member\u001b[39;00m\n",
      "\n",
      "    @property\n",
      "    \u001b[38;5;28;01mdef\u001b[39;00m func(self):\n",
      "        \u001b[38;5;28;01mreturn\u001b[39;00m _ffi_api.EnvFuncGetFunction(self)  \u001b[38;5;66;03m# type: ignore # pylint: disable=no-member\u001b[39;00m\n",
      "\n",
      "    @staticmethod\n",
      "    \u001b[38;5;28;01mdef\u001b[39;00m get(name):\n",
      "        \u001b[33m\"\"\"Get a static env function\u001b[39m\n",
      "\n",
      "\u001b[33m        Parameters\u001b[39m\n",
      "\u001b[33m        ----------\u001b[39m\n",
      "\u001b[33m        name : str\u001b[39m\n",
      "\u001b[33m            The name of the function.\u001b[39m\n",
      "\u001b[33m        \"\"\"\u001b[39m\n",
      "        \u001b[38;5;28;01mreturn\u001b[39;00m _ffi_api.EnvFuncGet(name)  \u001b[38;5;66;03m# type: ignore # pylint: disable=no-member\u001b[39;00m\n",
      "\u001b[31mFile:\u001b[39m           /media/pc/data/lxw/ai/tvm/python/tvm/ir/base.py\n",
      "\u001b[31mType:\u001b[39m           type\n",
      "\u001b[31mSubclasses:\u001b[39m     "
     ]
    }
   ],
   "source": [
    "tvm.ir.EnvFunc??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tvm.register_func(\"test.env_func\")\n",
    "def test(x):\n",
    "    return x + 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = tvm.get_global_func(\"test.env_func\")\n",
    "x = tvm.ir.EnvFunc.get(\"test.env_func\")\n",
    "assert x.name == \"test.env_func\"\n",
    "json_str = tvm.ir.save_json([x])\n",
    "y = tvm.ir.load_json(json_str)[0]\n",
    "assert y.name == x.name\n",
    "assert y(1) == 2\n",
    "assert y.func(1) == 2\n",
    "\n",
    "x = tvm.ir.make_node(\"attrs.TestAttrs\", name=\"xx\", padding=(3, 4), func=y)\n",
    "# assert x.name == \"xx\"\n",
    "assert x.padding[0].value == 3\n",
    "assert x.padding[1].value == 4\n",
    "assert x.axis == 10\n",
    "x = tvm.ir.load_json(tvm.ir.save_json(x))\n",
    "assert isinstance(x.func, tvm.ir.EnvFunc)\n",
    "assert x.func(10) == 11"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tvm.ir.base.EnvFunc,\n",
       " tvm.runtime.packed_func.PackedFunc,\n",
       " tvm.runtime.packed_func.PackedFunc)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(y), type(f), type(y.func)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "import numpy as np\n",
    "import pytest\n",
    "from tvm import te\n",
    "from tvm.ffi.access_path import AccessPath\n",
    "from tvm.script import tir as T, ir as I\n",
    "\n",
    "\n",
    "def consistent_equal(x, y, map_free_vars=False):\n",
    "    struct_equal0 = tvm.ir.structural_equal(x, y, map_free_vars)\n",
    "    struct_equal1 = tvm.ir.structural_equal(y, x, map_free_vars)\n",
    "\n",
    "    xhash = tvm.ir.structural_hash(x, map_free_vars)\n",
    "    yhash = tvm.ir.structural_hash(y, map_free_vars)\n",
    "\n",
    "    if struct_equal0 != struct_equal1:\n",
    "        raise ValueError(\n",
    "            \"Non-commutative {} vs {}, sequal0={}, sequal1={}\".format(\n",
    "                x, y, struct_equal0, struct_equal1\n",
    "            )\n",
    "        )\n",
    "\n",
    "    # NOTE: hash colision can happen but should be rare.\n",
    "    # we can confirm that hash colison doesn't happen for our testcases\n",
    "    if struct_equal0 != (xhash == yhash):\n",
    "        raise ValueError(\n",
    "            \"Inconsistent {} vs {}, sequal={}, xhash={}, yhash={}\".format(\n",
    "                x, y, struct_equal0, xhash, yhash\n",
    "            )\n",
    "        )\n",
    "    return struct_equal0\n",
    "@tvm.register_func(\"test.sequal.env_func\")\n",
    "def test(x):\n",
    "    return x + 1\n",
    "\n",
    "x = tvm.ir.EnvFunc.get(\"test.sequal.env_func\")\n",
    "y = tvm.ir.EnvFunc.get(\"test.sequal.env_func\")\n",
    "assert consistent_equal(y, x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py313",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
